from espdl import CocoDetector
from jpeg import Decoder, Encoder
from myufont import CustomBMFont
from machine import Pin, SDCard
import os

# Initialize SD card
sd = SDCard(slot=0, width=4, sck=43, cmd=44, data=(39, 40, 41, 42))
os.mount(sd, '/sd')

# Initialize decoders, encoders and detectors
decoder = Decoder()
encoder = Encoder(width=405, height=540, pixel_format="RGB888")
object_detector = CocoDetector(width=405, height=540)

# MSCOCO class labels
MSCOCO_CLASSES = [
    "person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
    "fire hydrant", "fire hose", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
    "elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
    "skis", "ski pole", "skateboard", "surfboard", "tennis racket", "bottle", "wine glass", "cup", "fork", "bowl",
    "banana", "apple", "sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake",
    "chair", "couch", "potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote",
    "keyboard", "cell phone", "microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase",
    "scissors", "teddy bear", "hair drier", "toothbrush"
]

# Load custom font
font = CustomBMFont('/sd/text_full_16px_2312.v3.bmf')

# Capture and process image
img = open("yolo.jpg", "rb").read()  # Capture original image (usually in JPEG format)
framebuffer = decoder.decode(img)  # Convert to RGB888 format
# Convert memoryview to bytearray for modification
framebuffer = bytearray(framebuffer)
# Run object detection
results = object_detector.run(framebuffer)

# Draw bounding box
def draw_rectangle(buffer, width, height, x, y, w, h, font, label, color=(255, 0, 0)):
    """
    Draw a rectangular border on an RGB888 format image buffer
    :param buffer: Image buffer
    :param width: Image width
    :param height: Image height
    :param x: X coordinate of the top-left corner of the border
    :param y: Y coordinate of the top-left corner of the border
    :param w: Border width
    :param h: Border height
    :param font: Font used for labeling
    :param label: Label text to display
    :param color: Border color (RGB format)
    """
    # Helper function: Set color of a single pixel
    def set_pixel(buffer, width, x, y, color):
        offset = (y * width + x) * 3
        buffer[offset] = color[0]  # R
        buffer[offset + 1] = color[1]  # G
        buffer[offset + 2] = color[2]  # B
    
    def is_chinese(ch):
        """Check if a character is a Chinese character"""
        if '\u4e00' <= ch <= '\u9fff' or \
           '\u3400' <= ch <= '\u4dbf' or \
           '\u20000' <= ch <= '\u2a6df':
            return True
        return False
    
    def text(font, text, x_start, y_start, color, spacing=0, line_spacing=0, max_width=width):
        font_size = font.font_size
        bytes_per_row = (font_size + 7) // 8  # Number of bytes per row
        x, y = x_start, y_start
        
        for char in text:
            # Handle newline characters
            if char == '\n':
                y += font_size + line_spacing
                x = x_start
                continue
            if char == '\r':
                x += 2 * font_size
                continue
            
            # Get character width (full width for Chinese, half width for ASCII)
            char_width = font_size if is_chinese(char) else font_size // 2
            
            # Check if line break is needed
            if max_width is not None and x + char_width > x_start + max_width:
                y += font_size + line_spacing
                x = x_start
            
            # Get character bitmap
            bitmap = font.get_char_bitmap(char)
            
            # Draw character
            for row in range(font_size):
                for col in range(char_width if not is_chinese(char) else font_size):
                    byte_idx = row * bytes_per_row + col // 8
                    bit_mask = 0x80 >> (col % 8)
                    
                    if byte_idx < len(bitmap) and (bitmap[byte_idx] & bit_mask):
                        set_pixel(framebuffer, max_width, x + col, y + row, color)
            
            # Move to next character position
            x += char_width + spacing
    
    # Draw top border
    for i in range(x, x + w):
        if 0 <= i < width and 0 <= y < height:
            set_pixel(buffer, width, i, y, color)

    # Draw bottom border
    for i in range(x, x + w):
        if 0 <= i < width and 0 <= y + h < height:
            set_pixel(buffer, width, i, y + h, color)

    # Draw left border
    for j in range(y, y + h):
        if 0 <= j < height and 0 <= x < width:
            set_pixel(buffer, width, x, j, color)

    # Draw right border
    for j in range(y, y + h):
        if 0 <= j < height and 0 <= x + w < width:
            set_pixel(buffer, width, x + w, j, color)
    
    # Draw label text
    text(font, label, x, y - 20, color)

# Draw borders on the image
for detection in results:
    #print(detection)
    x1, y1, x2, y2 = detection['box']
    label = f"{MSCOCO_CLASSES[detection['category']]}:{int(detection['score']*100)}%"
    draw_rectangle(framebuffer, 405, 540, x1, y1, x2 - x1, y2 - y1, font, label)  # Use red border
    print(label)

# Re-encode the image with borders into JPEG format and save
marked_img = encoder.encode(framebuffer)
with open("yolo_marked.jpg", "wb") as f:
    f.write(marked_img)